#!/bin/bash
# ========================================
# RedFire PostgreSQL用户初始化脚本
# 版本: v5.0.0 (2025-09-16)
# 功能: 创建业务用户和权限配置
# ========================================

set -e

echo "=========================================="
echo "开始初始化RedFire PostgreSQL用户和权限..."
echo "=========================================="

# 等待PostgreSQL完全启动
while ! pg_isready -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" -h localhost; do
    echo "等待PostgreSQL启动中..."
    sleep 2
done

# 连接到数据库并创建用户
psql -v ON_ERROR_STOP=1 --username "${POSTGRES_USER}" --dbname "${POSTGRES_DB}" <<-EOSQL

-- ========================================
-- 创建业务用户
-- ========================================

-- 创建只读用户 (查询权限)
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_user WHERE usename = 'redfire_readonly') THEN
        CREATE USER redfire_readonly WITH PASSWORD 'redfire_readonly_2025';
        RAISE NOTICE 'Created user: redfire_readonly';
    ELSE
        RAISE NOTICE 'User already exists: redfire_readonly';
    END IF;
END
\$\$;

-- 创建监控用户 (监控权限)
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_user WHERE usename = 'redfire_monitor') THEN
        CREATE USER redfire_monitor WITH PASSWORD 'redfire_monitor_2025';
        RAISE NOTICE 'Created user: redfire_monitor';
    ELSE
        RAISE NOTICE 'User already exists: redfire_monitor';
    END IF;
END
\$\$;

-- 创建备份用户 (备份权限)
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_user WHERE usename = 'redfire_backup') THEN
        CREATE USER redfire_backup WITH PASSWORD 'redfire_backup_2025';
        RAISE NOTICE 'Created user: redfire_backup';
    ELSE
        RAISE NOTICE 'User already exists: redfire_backup';
    END IF;
END
\$\$;

-- 创建复制用户 (主从复制权限)
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_user WHERE usename = 'redfire_repl') THEN
        CREATE USER redfire_repl WITH REPLICATION PASSWORD 'redfire_repl_2025';
        RAISE NOTICE 'Created user: redfire_repl';
    ELSE
        RAISE NOTICE 'User already exists: redfire_repl';
    END IF;
END
\$\$;

-- 创建数据同步用户 (ETL权限)
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_user WHERE usename = 'redfire_etl') THEN
        CREATE USER redfire_etl WITH PASSWORD 'redfire_etl_2025';
        RAISE NOTICE 'Created user: redfire_etl';
    ELSE
        RAISE NOTICE 'User already exists: redfire_etl';
    END IF;
END
\$\$;

-- 创建ML用户 (机器学习权限)
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_user WHERE usename = 'redfire_ml') THEN
        CREATE USER redfire_ml WITH PASSWORD 'redfire_ml_2025';
        RAISE NOTICE 'Created user: redfire_ml';
    ELSE
        RAISE NOTICE 'User already exists: redfire_ml';
    END IF;
END
\$\$;

-- 创建应用用户 (应用程序权限)
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_user WHERE usename = 'redfire_app') THEN
        CREATE USER redfire_app WITH PASSWORD 'redfire_app_2025';
        RAISE NOTICE 'Created user: redfire_app';
    ELSE
        RAISE NOTICE 'User already exists: redfire_app';
    END IF;
END
\$\$;

-- ========================================
-- 配置用户属性
-- ========================================

-- 设置连接限制
ALTER USER redfire_readonly CONNECTION LIMIT 50;
ALTER USER redfire_monitor CONNECTION LIMIT 10;
ALTER USER redfire_backup CONNECTION LIMIT 5;
ALTER USER redfire_etl CONNECTION LIMIT 20;
ALTER USER redfire_ml CONNECTION LIMIT 30;
ALTER USER redfire_app CONNECTION LIMIT 100;

-- 设置用户有效期 (可选)
-- ALTER USER redfire_readonly VALID UNTIL '2026-12-31';

-- ========================================
-- 创建角色和权限组
-- ========================================

-- 创建只读角色
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'readonly_role') THEN
        CREATE ROLE readonly_role;
        RAISE NOTICE 'Created role: readonly_role';
    END IF;
END
\$\$;

-- 创建分析师角色
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'analyst_role') THEN
        CREATE ROLE analyst_role;
        RAISE NOTICE 'Created role: analyst_role';
    END IF;
END
\$\$;

-- 创建数据工程师角色
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'data_engineer_role') THEN
        CREATE ROLE data_engineer_role;
        RAISE NOTICE 'Created role: data_engineer_role';
    END IF;
END
\$\$;

-- ========================================
-- 授予数据库级权限
-- ========================================

-- 只读用户权限
GRANT CONNECT ON DATABASE redfire_analytics TO redfire_readonly;
GRANT USAGE ON SCHEMA public TO redfire_readonly;
GRANT SELECT ON ALL TABLES IN SCHEMA public TO redfire_readonly;
GRANT SELECT ON ALL SEQUENCES IN SCHEMA public TO redfire_readonly;
GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA public TO redfire_readonly;

-- 监控用户权限
GRANT CONNECT ON DATABASE redfire_analytics TO redfire_monitor;
GRANT USAGE ON SCHEMA public TO redfire_monitor;
GRANT SELECT ON ALL TABLES IN SCHEMA public TO redfire_monitor;
GRANT pg_monitor TO redfire_monitor;

-- 备份用户权限
GRANT CONNECT ON DATABASE redfire_analytics TO redfire_backup;
GRANT USAGE ON SCHEMA public TO redfire_backup;
GRANT SELECT ON ALL TABLES IN SCHEMA public TO redfire_backup;

-- ETL用户权限
GRANT CONNECT ON DATABASE redfire_analytics TO redfire_etl;
GRANT USAGE ON SCHEMA public TO redfire_etl;
GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO redfire_etl;
GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO redfire_etl;
GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA public TO redfire_etl;

-- ML用户权限
GRANT CONNECT ON DATABASE redfire_analytics TO redfire_ml;
GRANT USAGE ON SCHEMA public TO redfire_ml;
GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO redfire_ml;
GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO redfire_ml;
GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA public TO redfire_ml;

-- 应用用户权限
GRANT CONNECT ON DATABASE redfire_analytics TO redfire_app;
GRANT USAGE ON SCHEMA public TO redfire_app;
GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO redfire_app;
GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO redfire_app;
GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA public TO redfire_app;

-- ========================================
-- 设置默认权限 (对未来创建的对象)
-- ========================================

-- 为只读用户设置默认权限
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO redfire_readonly;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON SEQUENCES TO redfire_readonly;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT EXECUTE ON FUNCTIONS TO redfire_readonly;

-- 为监控用户设置默认权限
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO redfire_monitor;

-- 为ETL用户设置默认权限
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO redfire_etl;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT USAGE, SELECT ON SEQUENCES TO redfire_etl;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT EXECUTE ON FUNCTIONS TO redfire_etl;

-- 为ML用户设置默认权限
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO redfire_ml;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT USAGE, SELECT ON SEQUENCES TO redfire_ml;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT EXECUTE ON FUNCTIONS TO redfire_ml;

-- 为应用用户设置默认权限
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO redfire_app;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT USAGE, SELECT ON SEQUENCES TO redfire_app;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT EXECUTE ON FUNCTIONS TO redfire_app;

-- ========================================
-- 授予角色权限
-- ========================================

-- 只读角色权限
GRANT readonly_role TO redfire_readonly;
GRANT USAGE ON SCHEMA public TO readonly_role;
GRANT SELECT ON ALL TABLES IN SCHEMA public TO readonly_role;

-- 分析师角色权限
GRANT readonly_role TO analyst_role;
GRANT analyst_role TO redfire_ml;
GRANT USAGE ON SCHEMA public TO analyst_role;
GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA public TO analyst_role;

-- 数据工程师角色权限
GRANT analyst_role TO data_engineer_role;
GRANT data_engineer_role TO redfire_etl;
GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO data_engineer_role;

-- ========================================
-- 创建用户管理函数
-- ========================================

-- 创建用户权限检查函数
CREATE OR REPLACE FUNCTION check_user_permissions(username text)
RETURNS TABLE(
    database_name text,
    schema_name text,
    object_name text,
    object_type text,
    privilege_type text
) AS \$\$
BEGIN
    RETURN QUERY
    SELECT 
        current_database()::text,
        n.nspname::text,
        c.relname::text,
        CASE c.relkind
            WHEN 'r' THEN 'table'
            WHEN 'v' THEN 'view'
            WHEN 'm' THEN 'materialized view'
            WHEN 'S' THEN 'sequence'
            WHEN 'f' THEN 'foreign table'
        END::text,
        p.privilege_type::text
    FROM information_schema.role_table_grants p
    JOIN pg_class c ON c.relname = p.table_name
    JOIN pg_namespace n ON n.nspname = p.table_schema
    WHERE p.grantee = username
        AND n.nspname = 'public'
    ORDER BY n.nspname, c.relname, p.privilege_type;
END;
\$\$ LANGUAGE plpgsql;

-- 创建用户连接统计函数
CREATE OR REPLACE FUNCTION get_user_connections()
RETURNS TABLE(
    username text,
    database_name text,
    connection_count bigint,
    active_queries bigint,
    idle_connections bigint
) AS \$\$
BEGIN
    RETURN QUERY
    SELECT 
        usename::text,
        datname::text,
        count(*)::bigint,
        count(*) FILTER (WHERE state = 'active')::bigint,
        count(*) FILTER (WHERE state = 'idle')::bigint
    FROM pg_stat_activity
    WHERE usename IS NOT NULL
    GROUP BY usename, datname
    ORDER BY count(*) DESC;
END;
\$\$ LANGUAGE plpgsql;

-- ========================================
-- 验证用户创建
-- ========================================

-- 显示所有用户
SELECT 
    usename AS username,
    usesuper AS is_superuser,
    usecreatedb AS can_create_db,
    useconnlimit AS connection_limit,
    valuntil AS valid_until
FROM pg_user 
WHERE usename LIKE 'redfire_%'
ORDER BY usename;

-- 显示所有角色
SELECT 
    rolname AS role_name,
    rolsuper AS is_superuser,
    rolinherit AS inherit_privileges,
    rolcreaterole AS can_create_roles,
    rolcreatedb AS can_create_db,
    rolconnlimit AS connection_limit
FROM pg_roles 
WHERE rolname LIKE '%role'
ORDER BY rolname;

-- 显示角色成员关系
SELECT 
    r.rolname AS role_name,
    m.rolname AS member_name
FROM pg_roles r
JOIN pg_auth_members am ON r.oid = am.roleid
JOIN pg_roles m ON am.member = m.oid
WHERE r.rolname LIKE '%role'
ORDER BY r.rolname, m.rolname;

EOSQL

echo "=========================================="
echo "RedFire PostgreSQL用户初始化完成！"
echo "创建的用户："
echo "  - redfire_readonly: 只读用户 (查询权限)"
echo "  - redfire_monitor: 监控用户 (监控权限)"
echo "  - redfire_backup: 备份用户 (备份权限)"
echo "  - redfire_repl: 复制用户 (复制权限)"
echo "  - redfire_etl: 数据同步用户 (ETL权限)"
echo "  - redfire_ml: 机器学习用户 (ML权限)"
echo "  - redfire_app: 应用用户 (应用权限)"
echo ""
echo "创建的角色："
echo "  - readonly_role: 只读角色"
echo "  - analyst_role: 分析师角色"
echo "  - data_engineer_role: 数据工程师角色"
echo ""
echo "创建的管理函数："
echo "  - check_user_permissions(): 检查用户权限"
echo "  - get_user_connections(): 获取用户连接统计"
echo "=========================================="
